/*
 * Copyright 2010-2012 Susanta Tewari. <freecode4susant@users.sourceforge.net>
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

package coalescent.is.cef.command;

import commons.util.TwoWayConverterS;
import commons.util.XMLUtil;
import commons.util.base.TwoWayConverter;

import org.dom4j.Document;
import org.dom4j.DocumentHelper;
import org.dom4j.Element;

import static commons.util.ExceptionUtil.throwArgEx;

import java.io.File;
import java.io.IOException;
import java.io.Serializable;
import java.util.Arrays;
import java.util.List;
import java.util.Objects;

import static java.lang.Integer.parseInt;

/**
 * Specifies and implements the load-store of the execution state of {@code AppNoteIS_Cmd} so that
 * partial results from multiple executions on the same or different machines can be combined.
 *
 * State includes the input of the command as well.
 *
 *
 *
 * @author Susanta Tewari
 * @version version
 * @history Created on 12/13/13.
 * @since since
 */
public final class AppNoteIS_CmdState {

    /** 2Way converters for double/int arrays */
    private static final TwoWayConverter<String, int[]> cnv2_i    = TwoWayConverterS.ofInt(",");
    private static final TwoWayConverter<String, double[]> cnv2_d = TwoWayConverterS.ofDouble(",");
    private static final TwoWayConverter<String, List<String>> cnvS =
        TwoWayConverterS.of_String_List(",", TwoWayConverterS.ofString());

    /** command input parameters */
    private final double[] thetas;
    private final int[] nVals;
    private final int scaledSamplerIndex;
    private final String[] samplers;
    private int sizePerOrder   = 100;
    private int cellSampleSize = 3;
    private int samplerCount   = 3;

    /** command output */
    private final Output[][][] gridOutputBySize;
    private final Output[][][] gridOutputByTime;
    private String measureKey;

    public static AppNoteIS_CmdState load(String statefulPath) throws Exception {

        Element root_e = XMLUtil.getRootElement(new File(statefulPath), false); //J-
        final double[] thetas          = cnv2_d.fwd(root_e.elementText("Y_mutationRates").trim());
        final int[] nVals              = cnv2_i.fwd((root_e.elementText("X_dataSampleSizes").trim()));
        final int samplerCount         = parseInt(root_e.elementText("samplerCount"));
        final String[] samplerLabels   = cnvS.fwd(root_e.elementText("samplerLabels").trim()).toArray(new String[0]);
        final int scaledSamplerIndex   = parseInt(root_e.elementText("scaledSamplerIndex"));

        final AppNoteIS_CmdState state = new AppNoteIS_CmdState(thetas, nVals, samplerLabels, scaledSamplerIndex);

        final List<Element> rows       = root_e.element("gridOutput").elements("row");//J+


        for (Element row_e : rows) {

            final int rowIndex       = parseInt(row_e.attributeValue("index"));
            final List<Element> colS = row_e.elements("col");

            for (Element col_e : colS) {

                final int colIndex = parseInt(col_e.attributeValue("index"));


                // size
                final Element bySize           = col_e.element("bySize");
                final int size                 = Integer.parseInt(bySize.attributeValue("size"));
                final List<Element> dataS_list = bySize.elements("data");

                for (int k = 0; k < dataS_list.size(); k++) {

                    double[] mean   = cnv2_d.fwd(dataS_list.get(k).elementText("mean").trim());
                    double[] se     = cnv2_d.fwd(dataS_list.get(k).elementText("se").trim());
                    double[] ess    = cnv2_d.fwd(dataS_list.get(k).elementText("ess").trim());
                    int[] seconds   = cnv2_i.fwd(dataS_list.get(k).elementText("seconds").trim());
                    int[] sizeArray = new int[samplerCount];

                    Arrays.fill(sizeArray, size);

                    final Output output = new Output(mean, se, ess, sizeArray, seconds);

                    state.addOutputBySize(output, rowIndex, colIndex, k);
                }


                // time
                final Element byTime           = col_e.element("byTime");
                final int seconds              = Integer.parseInt(byTime.attributeValue("seconds"));
                final List<Element> dataT_list = byTime.elements("data");

                for (int k = 0; k < dataT_list.size(); k++) {

                    double[] mean      = cnv2_d.fwd(dataT_list.get(k).elementText("mean").trim());
                    double[] se        = cnv2_d.fwd(dataT_list.get(k).elementText("se").trim());
                    double[] ess       = cnv2_d.fwd(dataT_list.get(k).elementText("ess").trim());
                    int[] sizeArray    = cnv2_i.fwd(dataT_list.get(k).elementText("size").trim());
                    int[] secondsArray = new int[samplerCount];

                    Arrays.fill(secondsArray, seconds);

                    final Output output = new Output(mean, se, ess, sizeArray, secondsArray);

                    state.addOutputByTime(output, rowIndex, colIndex, k);
                }
            }
        }

        final int sizePerOrder   = parseInt(root_e.elementText("simSizePerOrder"));
        final int cellSampleSize = parseInt(root_e.elementText("XY_gridCellSampleCount"));

        state.setSizePerOrder(sizePerOrder);
        state.setCellSampleSize(cellSampleSize);
        state.setSamplerCount(samplerCount);

        return state;
    }

    public static void store(AppNoteIS_CmdState state, String statefulPath) throws IOException {

        Document document    = DocumentHelper.createDocument();
        final Element root_e = document.addElement("root");

        root_e.addAttribute("cmd", "coalescent.is.cef.command$AppNoteIS_Cmd");
        root_e.addElement("simSizePerOrder").addText(state.getSizePerOrder() + "");
        root_e.addElement("XY_gridCellSampleCount").addText(state.getCellSampleSize() + "");
        root_e.addElement("samplerCount").addText(state.getSamplerCount() + "");
        root_e.addElement("samplerLabels").addText(cnvS.bck(Arrays.asList(state.getSamplers())));
        root_e.addElement("scaledSamplerIndex").addText(state.getScaledSamplerIndex() + "");


        // thetas & nVals
        final double[] thetas = state.getThetas();

        root_e.addElement("Y_mutationRates").addText(cnv2_d.bck(thetas));

        final int[] nVals = state.getNVals();

        root_e.addElement("X_dataSampleSizes").addText(cnv2_i.bck(nVals));


        // grid
        final Element grid_e = root_e.addElement("gridOutput");

        for (int i = 0; i < thetas.length; i++) {

            final Element row_e = grid_e.addElement("row");

            row_e.addAttribute("index", "" + i);

            for (int j = 0; j < nVals.length; j++) {

                if (!state.containsOutputBySize(i, j, 0)) continue;

                final Element col_e = row_e.addElement("col");

                col_e.addAttribute("index", "" + j);

                final Element bySize_e     = col_e.addElement("bySize");
                final Element byTime_e     = col_e.addElement("byTime");
                final Element dataBySize_e = bySize_e.addElement("data");
                final Element dataByTime_e = byTime_e.addElement("data");

                for (int k = 0; k < state.getCellSampleSize(); k++) {

                    final Output oSize = state.getOutputBySize(i, j, k);

                    dataBySize_e.addElement("mean").addText(cnv2_d.bck(oSize.getMean()));
                    dataBySize_e.addElement("se").addText(cnv2_d.bck(oSize.getSe()));
                    dataBySize_e.addElement("ess").addText(cnv2_d.bck(oSize.getEss()));
                    dataBySize_e.addElement("size").addText(cnv2_i.bck(oSize.getSize()));
                    dataBySize_e.addElement("seconds").addText(cnv2_i.bck(oSize.getSeconds()));

                    final Output oTime = state.getOutputByTime(i, j, k);

                    dataByTime_e.addElement("mean").addText(cnv2_d.bck(oTime.getMean()));
                    dataByTime_e.addElement("se").addText(cnv2_d.bck(oTime.getSe()));
                    dataByTime_e.addElement("ess").addText(cnv2_d.bck(oTime.getEss()));
                    dataByTime_e.addElement("size").addText(cnv2_i.bck(oTime.getSize()));
                    dataByTime_e.addElement("seconds").addText(cnv2_i.bck(oTime.getSeconds()));
                }
            }
        }

        XMLUtil.write(document, statefulPath);
    }

    public AppNoteIS_CmdState(double[] thetas, int[] nVals, String[] samplers,
                              int scaledSamplerIndex) {

        this.thetas             = thetas;
        this.nVals              = nVals;
        this.samplers           = samplers;
        this.scaledSamplerIndex = scaledSamplerIndex;
        this.gridOutputBySize   = new Output[thetas.length][nVals.length][cellSampleSize];
        this.gridOutputByTime   = new Output[thetas.length][nVals.length][cellSampleSize];
    }

    /**
     * @param output output to be added
     * @param rowIndex row index in the grid;must be non-negative and less than size of
     * {@code getThetas()}
     * @param colIndex column index in the grid; must be non-negative and less than size of
     * {@code getNVals()}
     * @throws NullPointerException if {@code output} is {@code null}
     * @throws IllegalArgumentException if arguments are invalid
     */
    public void addOutputBySize(Output output, int rowIndex, int colIndex, int zIndex) {

        Objects.requireNonNull(output);
        checkRowColIndices(rowIndex, colIndex, zIndex);

        gridOutputBySize[rowIndex][colIndex][zIndex] = output;
    }

    public void addOutputByTime(Output output, int rowIndex, int colIndex, int zIndex) {

        Objects.requireNonNull(output);
        checkRowColIndices(rowIndex, colIndex, zIndex);

        gridOutputByTime[rowIndex][colIndex][zIndex] = output;
    }

    /**
     * @param rowIndex row index in the grid;must be non-negative and less than size of
     * {@code getThetas()}
     * @param colIndex column index in the grid; must be non-negative and less than size of
     * {@code getNVals()}
     * @return IllegalArgumentException if arguments are invalid
     */
    public boolean containsOutputBySize(int rowIndex, int colIndex, int zIndex) {

        checkRowColIndices(rowIndex, colIndex, zIndex);

        return gridOutputBySize[rowIndex][colIndex][zIndex] != null;
    }

    public boolean containsOutputByTime(int rowIndex, int colIndex, int zIndex) {

        checkRowColIndices(rowIndex, colIndex, zIndex);

        return gridOutputByTime[rowIndex][colIndex][zIndex] != null;
    }

    /**
     * @param rowIndex row index in the grid;must be non-negative and less than size of
     * {@code getThetas()}
     * @param colIndex column index in the grid; must be non-negative and less than size of
     * {@code getNVals()}
     * @return IllegalArgumentException if arguments are invalid
     */
    public Output getOutputBySize(int rowIndex, int colIndex, int zIndex) {

        checkRowColIndices(rowIndex, colIndex, zIndex);

        return gridOutputBySize[rowIndex][colIndex][zIndex];
    }

    public Output getOutputByTime(int rowIndex, int colIndex, int zIndex) {

        checkRowColIndices(rowIndex, colIndex, zIndex);

        return gridOutputByTime[rowIndex][colIndex][zIndex];
    }

    private void checkRowColIndices(int rowIndex, int colIndex, int zIndex) {

        throwArgEx((rowIndex < 0) || (rowIndex >= thetas.length), "invalid row index");
        throwArgEx((colIndex < 0) || (colIndex >= nVals.length), "invalid col index");
        throwArgEx((zIndex < 0) || (zIndex >= cellSampleSize), "invalid z-index");
    }


    // <editor-fold desc="Getters & Setters">
    public double[] getThetas() {
        return thetas;
    }

    public String[] getSamplers() {
        return samplers;
    }

    public int[] getNVals() {
        return nVals;
    }

    public int getSizePerOrder() {
        return sizePerOrder;
    }

    public int getCellSampleSize() {
        return cellSampleSize;
    }

    public int getSamplerCount() {
        return samplerCount;
    }

    public int getScaledSamplerIndex() {
        return scaledSamplerIndex;
    }

    public void setSizePerOrder(int sizePerOrder) {
        this.sizePerOrder = sizePerOrder;
    }

    public void setCellSampleSize(int cellSampleSize) {
        this.cellSampleSize = cellSampleSize;
    }

    public void setSamplerCount(int samplerCount) {
        this.samplerCount = samplerCount;
    }    // </editor-fold>

    public String getMeasureKey() {
        return measureKey;
    }

    /**
     *
     */
    public static final class Output implements Serializable {

        private final double[] mean;
        private final double[] se;
        private final double[] ess;
        private final int[] size;
        private final int[] seconds;

        /**
         * 3 data points and ESS scaled to SD for EGT and HUW (in order)
         * @param size
         * @param seconds
         */
        public Output(double[] mean, double[] se, double[] ess, int[] size, int[] seconds) {

            Objects.requireNonNull(mean);
            Objects.requireNonNull(se);
            Objects.requireNonNull(ess);
            Objects.requireNonNull(size);
            Objects.requireNonNull(seconds);

            this.size    = size;
            this.seconds = seconds;
            this.mean    = mean;
            this.se      = se;
            this.ess     = ess;
        }

        public double[] getMean() {
            return mean;
        }

        public double[] getSe() {
            return se;
        }

        public double[] getEss() {
            return ess;
        }

        public double[] getMeasure(String key) {

            switch (key) {

                case "se" :
                    return getSe();

                case "ess" :
                    return getEss();

                default :
                    throw new IllegalArgumentException("Invalid key");
            }
        }

        public int[] getSize() {
            return size;
        }

        public int[] getSeconds() {
            return seconds;
        }
    }
}
